
#define BOOST_NO_CXX11_SCOPED_ENUMS
#include "yolov3_clip_detector.h"
#include <algorithm>
#include <dirent.h>
#include <math.h>
#include <cstdlib>
#include <chrono>
#include <unistd.h>
#include <boost/filesystem.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

#include <NvOnnxParser.h>


namespace optical_clip_detector
{

  void createDirs(const std::string &dir_name)
  {
    if (!boost::filesystem::exists(dir_name))
    {
      try
      { boost::filesystem::create_directories(dir_name); }
      catch (...)
      {
        std::cout<<"Fail to create directory: "<<dir_name.c_str()<<", exits"<<std::endl;
        exit(-1);
      }
    }
  }

  /**
 * @brief The interface class of imagebind clip classify
 * 
 */
  Yolov3CLIPDetector::Yolov3CLIPDetector()
  {
    std::cout<<"-- Creat instence of Resource Yolov3CLIPDetector"<<std::endl;
  }

  Yolov3CLIPDetector::~Yolov3CLIPDetector()
  {
      _context->destroy();
      _engine->destroy();
      runtime->destroy();
      for (auto &binding:_bindings)
      {
        cudaCheck(cudaFree(binding));
      }
      std::cout<<"-- Destroy instance of Resource Yolov3CLIPDetector"<<std::endl;
  }

  int getElementSize(nvinfer1::DataType &type) {
      switch(type)
      {
       case nvinfer1::DataType::kFLOAT :   return 4;
       case nvinfer1::DataType::kHALF :   return 2;
       case nvinfer1::DataType::kINT8 :   return 1;
       case nvinfer1::DataType::kINT32 :   return 4;
       case nvinfer1::DataType::kBOOL :   return 4;
      }
  }
  
  int volume(const nvinfer1::Dims& d) {
      return std::accumulate(d.d, d.d + d.nbDims, 1, std::multiplies<int>());
  }
   /**
 * @brief prepare session option for cpu
 * @return bool 
 * 
 */
  nvinfer1::ICudaEngine* Yolov3CLIPDetector::getEngine(const std::string& engine) {

    std::ifstream engineFile(engine, std::ios::binary);
    if (!engineFile)
    {
        std::cout<<"Error opening engine file: "<<engine.c_str()<<std::endl;
        return nullptr;
    }

    engineFile.seekg(0, engineFile.end);
    long int fsize = engineFile.tellg();
    engineFile.seekg(0, engineFile.beg);

    std::vector<char> engineData(fsize);
    engineFile.read(engineData.data(), fsize);
    if (!engineFile)
    {
        std::cout<<"Error loading engine file: "<<engine<<std::endl;
        return nullptr;
    }

    runtime = nvinfer1::createInferRuntime(logger);

    return runtime->deserializeCudaEngine(engineData.data(), fsize, nullptr);
  }

  bool Yolov3CLIPDetector::initEngine(const std::string& model_path){
      // 2. create asynchronous cuda stream
    cudaCheck(cudaStreamCreate(&_stream));

    // you must reset context first, then you can reset the engine.
    // otherwise the context destructor will fail, a segment fault happens.
    //_context.reset();
    // cudaCheck(cudaSetDevice(_gpu_id));
    _engine = getEngine(model_path);
    if (!_engine)
    {
      std::cout<<"Engine creation failed"<<std::endl;
      return false;
    }

    // 4. create context
    _context = _engine->createExecutionContext();

    // 5. set input/output buffer
    auto num_binds = _engine->getNbBindings();
    _bindings.resize(num_binds);
    _buffer_size.resize(num_binds);
    for (int i = 0; i < num_binds; i++)
    {
        nvinfer1::Dims dims = _engine->getBindingDimensions(i);    // (3, 224, 224)  (1024)
        if (_engine->bindingIsInput(i))
        {
            const bool isDynamicInput = std::any_of(dims.d, dims.d + dims.nbDims, [](int dim){ return dim == -1; });
            // set dynamic dimensions to 1.
            if (isDynamicInput)
            {
                std::vector<int> staticDims;
                staticDims.resize(dims.nbDims);
                std::transform(dims.d, dims.d + dims.nbDims, staticDims.begin(),
                    [&](int dim) { return dim >= 0 ? dim : 1; });
                nvinfer1::Dims new_dims{static_cast<int>(staticDims.size()), {}};
                std::copy_n(staticDims.begin(), new_dims.nbDims, std::begin(new_dims.d));
                _context->setBindingDimensions(i, new_dims);
            }
        }
        dims = _context->getBindingDimensions(i);
        nvinfer1::DataType dtype = _engine->getBindingDataType(i); // 0, 0 也就是两个都是kFloat类型
        int64_t total_size = volume(dims) * 1 * getElementSize(dtype);
        _buffer_size[i] = total_size;
        cudaCheck(cudaMalloc(&_bindings[i], total_size));
    }
    _embeddings.resize(_buffer_size[2] / sizeof(float));
    _embeddings1.resize(_buffer_size[3] / sizeof(float));
    return true;
  }

 /**
 * @brief The interface class of init
 * @param resource_id
 * @return true or false
 * 
 */  
 bool Yolov3CLIPDetector::Initialize()
  {
    bool flag = false;
    // init imagebind config


    const std::string model_path =  "../merge.engine";

    if (cls_names.size() >0){
      cls_names.clear();
    }
  
    target_names = cls_names;
    
    text_emdeding_path = "/home/xjl/ccode/imagebind_test_1025/text_vectors";
    imagebind_class_names = {"Background", "Warplane", "Himars_carrier","Himars_launcher","Hummer"};
    std::cout<<"text_embeding_path = "<< text_emdeding_path<<std::endl;
    std::vector<std::string> text_objects = {"Background","Warplane"};
    cls_names.insert(cls_names.end(),text_objects.begin(), text_objects.end());

    img_padding = 20;
    relearning_mode = 0;
    width_height_ratio = 4;
    mode2_update_threshold = 0;
    mode2_remove_threshold = 0;
    mode1_update_threshold = 0;
    mode1_remove_threshold = 0;
    temperature = 0.1;
    modify_mode = 1;
    max_target_num = 2;
    imagebind_detect_order = 0;
    save_img_flag = 0;
    
       
    flag = initEngine(model_path);
    if (!flag){
      std::cout<<"init model false"<<std::endl;
      return false;
    }
    
    if (_known_embeddings.size() > 0){
      _known_embeddings.clear();
    }
    load_pre_text_embeddings(text_emdeding_path, _known_embeddings, cls_names);
    std::cout<<"known_embeddings size = "<< _known_embeddings.size()<<std::endl;
    return true;
  }

// /**
//  * @brief The interface class of setParams
//  * @param aConfidenceThres
//  * @param aWidth
//  * @param aHeight
//  * @return void
//  * 
//  */ 
//   void Yolov3CLIPDetector::setParams(float aConfidenceThres,                           // Confidence threshold
//                                       std::unordered_set<std::string> aTargetClassname, // Classes need to be detected
//                                       std::string aModelFilePath,                       // Model file path
//                                       size_t aWidth,                                    // Block width to DNN
//                                       size_t aHeight)
//   {
//     std::cout<<"SetParameter of YoloV3GPUDetector..."<<std::endl;
//     _YoloV3GPUDetector->setParams(aConfidenceThres, aTargetClassname, aModelFilePath, aWidth, aHeight);
//   }




/**
 * @brief The interface class of inference
 * @param pFrame
 * @param pRoi
 * @param pName
 * @param pConf
 * @return true or false
 * 
 */


  bool Yolov3CLIPDetector::inference(cv::Mat &visible_pFrame,
                                  std::vector<cv::Rect2d> &pRoi,
                                  std::vector<std::string> &pName,
                                  std::vector<float> &pConf
                                  )
  {
    bool detected = true;
    pRoi.push_back(cv::Rect(40, 40, 80, 80));
    pName.push_back("Warplane");
    pConf.push_back(0.8);
    pRoi.push_back(cv::Rect(40, 40, 80, 80));
    pName.push_back("Warplane");
    pConf.push_back(0.8);
    
    std::cout<<"pName = "<< pName.size()<< " pConf = "<<pConf.size()<<std::endl;
    for (int i=0; i<pName.size(); i++)
    {
      std::cout<<"i = "<<i<<" pName = "<<pName[i]<<" pConf = "<<pConf[i]<<std::endl;
    }
    // maximum processed target number
    int target_num = std::min(max_target_num, int(pRoi.size()));
    std::cout<<"target_num = "<< target_num<<std::endl;
    int valid_targets = 0;
    int modify_count = 0;
    std::cout<<"size = "<< _known_embeddings.size()<<std::endl;
    std::vector<int> rm_index;
    if (_known_embeddings.size() > 1)
    {
      for (int i = 0; i < target_num; i++)
      {
        int index;
        if(imagebind_detect_order == 1)
        {
          index = pRoi.size() - i - 1;
          std::cout<<"ascend order: index = "<<index<<"-------------------------------------------------"<<std::endl;
        } else 
        {
          index = i;
          std::cout<<"descend order: index = "<<index<<"-------------------------------------------------"<<std::endl;
        }
        cv::Rect2d Roi = pRoi[index];
        std::string yolo_name = pName[index];
        float yolo_conf = pConf[index];
        auto it = std::find(imagebind_class_names.begin(),imagebind_class_names.end(),yolo_name);
        if(it ==imagebind_class_names.end()){
          std::cout<<"imagebind does not support predicting: "<< yolo_name.c_str()<<std::endl;
          continue;
        }
        float radio1 =  Roi.width / (Roi.height + std::numeric_limits<float>::epsilon());
        float radio2 =  Roi.height / (Roi.width + std::numeric_limits<float>::epsilon());
        float max_ratio = std::max(radio1, radio2);
        if (max_ratio > width_height_ratio){
            modify_count++;
            rm_index.push_back(index);
            std::cout<<"yolo predict object error"<<std::endl;
            continue;
        }
        
        int x1 = std::max(0, int(Roi.x - img_padding / 2));
        int y1 = std::max(0, int(Roi.y - img_padding / 2));
        int pwidth = std::min(int(x1 + Roi.width + img_padding), int(visible_pFrame.cols -1)) - x1;
        int pheight = std::min(int(y1 + Roi.height + img_padding), int(visible_pFrame.rows -1)) - y1;

        cv::Rect largeRect(x1, y1, pwidth, pheight);
        //cv::Mat image_roi = visible_pFrame(largeRect);  
        cv::Mat image_roi = cv::imread("../dataset/1.jpg");
        std::vector<float> input_data;
        std::vector<float> input_data2;

        for(int i=0; i<4; i++){
          std::cout<<"input and out: "<<_buffer_size[i]<<std::endl;
        }
        std::cout<<"_clip_Width: "<<_clip_Width<<", _clip_Height: "<<_clip_Height<<std::endl;
        image_preprocess(image_roi, _clip_Width, _clip_Height, input_data, means, stds);
        image_preprocess(image_roi, _clip_Width, _clip_Height, input_data2, means, stds);
        cudaCheck(cudaMemcpyAsync(_bindings[0], input_data.data(), _buffer_size[0], cudaMemcpyHostToDevice, _stream));
        cudaCheck(cudaMemcpyAsync(_bindings[1], input_data2.data(), _buffer_size[1], cudaMemcpyHostToDevice, _stream));
        _context->enqueueV2(&_bindings[0], _stream, nullptr);
      
        cudaCheck(cudaMemcpyAsync(_embeddings.data(), _bindings[2], _buffer_size[2], cudaMemcpyDeviceToHost, _stream));
        cudaCheck(cudaMemcpyAsync(_embeddings1.data(), _bindings[3], _buffer_size[3], cudaMemcpyDeviceToHost, _stream));
        cudaCheck(cudaStreamSynchronize(_stream));

        std::cout<<"-------------"<<std::endl;
        for (int i=0; i<15; i++){
          std::cout<<input_data[i]<<" ";

        }
        std::cout<<" "<<std::endl;
        std::cout<<"----------------"<<std::endl;
        for (int i=0; i<6; i++){
          std::cout<<"_embeddings val: "<<_embeddings[i]<<std::endl;
        }

        for (int i=0; i<6; i++){
          std::cout<<"_embeddings1 val: "<<_embeddings1[i]<<std::endl;
        }
        valid_targets += 1;
        std::vector<float> imagebind_conf;
        std::vector<std::string> imagebind_names;
        for (auto &known_embedding:_known_embeddings)
        {
          float cls_in_conf=-1000;
          std::string cls_in_name;
          for (auto &it : known_embedding)
          {
            float val = std::inner_product(_embeddings.begin(), _embeddings.end(), it.second.begin(), 0.0);
            if (val > cls_in_conf)
            {
              cls_in_conf = val;
            }
            cls_in_name = it.first.substr(0, it.first.find_last_of("_"));
          }
          imagebind_conf.push_back(cls_in_conf);
          imagebind_names.push_back(cls_in_name);
        }

        std::vector<float> softmax_conf = softmax(imagebind_conf, temperature);
        // cal max value and index
        std::vector<float>::iterator biggest = std::max_element(softmax_conf.begin(), softmax_conf.end());
        float max_value = *biggest;
        int max_value_index = std::distance(softmax_conf.begin(), biggest);
        std::string imagebindPredName = imagebind_names[max_value_index];

        std::cout<<"IMAGEBIND: Current box YOLO category: "<<yolo_name<<", ImageBind category: "<<imagebindPredName<<std::endl;
        for (int j = 0; j < softmax_conf.size(); j++)
        {
          std::cout<<"IMAGEBIND: index = "<<j<<" name: "<<imagebind_names[j]<<" softmax_conf: "<<softmax_conf[j]<<std::endl;
        }

        if (modify_mode == 1)
        {   
          auto its = std::find(target_names.begin(),target_names.end(),imagebindPredName);
          bool valild_object_flag = its ==target_names.end();
          if (valild_object_flag)
          { 
            if (max_value > mode1_remove_threshold){
              modify_count++;
              rm_index.push_back(index);
              std::cout<<"IMAGEBIND: TARGET REMOVED!!! Current detection box category: "<<yolo_name<<", ImageBind Max category confidence is: "<<max_value<<" > "<<mode1_remove_threshold<<", this box should be False Positive (background)"
                            <<std::endl;
            }else{
              std::cout<<"IMAGEBIND: TARGET KEEP!!! Current detection box category: "<<yolo_name<<", ImageBind predicted category:"<<imagebindPredName<<", but Max category confidence is: "<<max_value<<" < "<<mode1_remove_threshold<<", so keep the object"<<std::endl;
            }

          }
          else
          {
              if (max_value > mode1_update_threshold){
                modify_count++;
                pName[index] = imagebind_names[max_value_index];
                pConf[index] = max_value;
                std::cout<<"IMAGEBIND: TARGET MODIFIED!!! Current detection box category: "<<yolo_name<<", ImageBind category: "<<imagebindPredName<<", ImageBind confidence is "<<max_value<<" > "<<mode1_update_threshold<<", modifying detection category to "<<imagebindPredName<<std::endl;
                                
              }else{
                 std::cout<<"IMAGEBIND: TARGET KEEP!!! Current detection box category: "<<yolo_name<<", ImageBind predicted category:"<<imagebindPredName<<", but Max category confidence is: "<<max_value<<"  < "<<mode1_update_threshold<<", so keep the object";
              }
          }

        } else
        {
          std::string imgbind_name = imagebind_names[max_value_index];
          if (imgbind_name != yolo_name){
            if ((max_value < mode2_remove_threshold) || imagebindPredName == "Background" || imagebindPredName == "Others"){
              modify_count++;
              rm_index.push_back(index);
              std::cout<<"IMAGEBIND: TARGET REMOVED!!! Current detection box category: "<<yolo_name<<", ImageBind Max category confidence is: "<<max_value<<" > "<<mode2_remove_threshold<<", this box should be False Positive (background)"
                            <<std::endl;

            }else if(max_value > mode2_update_threshold){
                modify_count++;
                pName[index] = imagebind_names[max_value_index];
                pConf[index] = max_value;
                std::cout<<"IMAGEBIND: TARGET MODIFIED!!! Current detection box category: "<<yolo_name<<", ImageBind category: "<<imagebindPredName<<", ImageBind confidence is "<<max_value<<" > "<<mode2_update_threshold<<", modifying detection category to "<<imagebindPredName<<std::endl;
            }else {
                std::cout<<"IMAGEBIND: TARGET KEEP!!! Current detection box category: "<<yolo_name<<", ImageBind predicted category:"<<imagebindPredName<<", but Max category confidence is: "<<max_value<<"  < "<<mode1_update_threshold<<", so keep the object";
            }

          }
          else {
              std::cout<<"TARGET KEPT!!!Current detection box category: "<<yolo_name<<", ImageBind category: "<<imagebindPredName<<std::endl;
            }
           
        }
      }
    } else 
    {
      std::cout<<"DO NOT have enough pretrained text embeddings, ImageBind no inference"<<std::endl;
    }

    for (int i=valid_targets; i< max_target_num; i++){
      std::cout<<"Run a empty img"<<std::endl;
      std::vector<float>input_vectors_rand(_clip_Width * _clip_Height * 3, 0);
      cudaCheck(cudaMemcpyAsync(_bindings[0], input_vectors_rand.data(), _buffer_size[0], cudaMemcpyHostToDevice, _stream));
      _context->enqueueV2(&_bindings[0], _stream, nullptr);
      cudaCheck(cudaMemcpyAsync(_embeddings.data(), _bindings[1], _buffer_size[1], cudaMemcpyDeviceToHost, _stream));
      cudaCheck(cudaStreamSynchronize(_stream));
    }
    
    // remove error Roi
    if (rm_index.size() > 0)
    {
      for (int i = 0; i < rm_index.size(); i++)
      {
        if (imagebind_detect_order == 0)
        {
          pRoi.erase(pRoi.begin()+rm_index[i]-i);
          pName.erase(pName.begin()+rm_index[i]-i);
          pConf.erase(pConf.begin()+rm_index[i]-i);
        } else
        {
          pRoi.erase(pRoi.begin()+rm_index[i]);
          pName.erase(pName.begin()+rm_index[i]);
          pConf.erase(pConf.begin()+rm_index[i]);
        }
      }
    }

    std::cout<<"IMAGEBIND: ImageBind modified boxes numbers: "<< modify_count<<std::endl;
    if (pRoi.size() == 0 && detected == true)
    {
      detected = false;
    }
    return detected;
  }
  

  bool image_preprocess(const cv::Mat &src_image, int width, int height, std::vector<float> &input_vectors, std::vector<float>&means, std::vector<float>&stds) {
    cv::Mat src_temp, dst_temp, dst_image;
    // 1. resize by CUBIC mode
    cv::resize(src_image, src_temp, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
    src_temp.convertTo(dst_temp, CV_32F);
    dst_temp.copyTo(dst_image);

    // 2. convert BGR to RGB
    cv::cvtColor(dst_image, dst_image, cv::COLOR_BGR2RGB);
    for (int i = 0; i < dst_image.rows; i++) {
        for (int j = 0; j < dst_image.cols; j++) {
            cv::Vec3f buf = dst_image.at<cv::Vec3f>(i,j);
            // printf("buf %f %f %f\n", buf[0], buf[1], buf[2]);
            buf[0] = ((buf[0] / 255.0f) - means[0]) / stds[0];
            buf[1] = ((buf[1] / 255.0f) - means[1]) / stds[1];
            buf[2] = ((buf[2] / 255.0f) - means[2]) / stds[2];
            dst_image.at<cv::Vec3f>(i,j) = buf;
            // printf("buf %f %f %f\n", buf[0], buf[1], buf[2]);
        }
    }

    // 4.convert mat to vector
    cv2vec(dst_image, input_vectors);

    return true;
  }

  bool cv2vec(const cv::Mat src_image, std::vector<float> &input_vectors) {
      cv::Mat channels[3]; //借用来进行HWC->CHW
      cv::split(src_image, channels);
      for (int i = 0; i < src_image.channels(); i++)  // HWC->CHW
      {
          std::vector<float> data = std::vector<float>(channels[i].reshape(1, src_image.cols * src_image.rows));
          input_vectors.insert(input_vectors.end(), data.begin(), data.end());
      }
      
      return true;
  }

  std::map<std::string ,std::vector<float>> load_pre_text_file(std::vector<std::string> &files) {
      std::map<std::string ,std::vector<float>> known_embedding;
      for (auto &file : files)
      {
          size_t pos1 = file.find_last_of('/');
          std::string folder = file.substr(0, pos1);
          std::string cls = folder.substr(folder.find_last_of('/')+1, folder.length());
          size_t pos2 = file.find('.');
          std::string index = file.substr(pos1+1, pos2-pos1-1);
          std::ifstream in_file(file);
          std::vector <float> record;
          while (in_file)
          {
              std::string s;
              if (!getline(in_file, s)) break;
              std::istringstream ss(s);
              while (ss)
              {
                  std::string s;
                  if (!getline(ss, s, ',')) break;
                  record.push_back(atof(s.c_str()));
              }
          }

          known_embedding[cls+"_"+index].assign(record.begin(), record.end());
      }
      return known_embedding;
  }

  bool load_pre_text_embeddings(std::string path, std::vector<std::map<std::string ,std::vector<float>>> &known_embeddings, std::vector<std::string> cls_names) {
    //file handle
    DIR *d = opendir(path.c_str());
    //file path
    std::vector<std::string> files;
    std::string p;
    if (d != NULL)
    {
      //if folder, iter,
      //else added to vector
      struct dirent *dt = NULL;
      while (dt = readdir(d))
      {
        if (strcmp(dt->d_name, ".") != 0 && strcmp(dt->d_name, "..") != 0)
        {
          if (dt->d_type == 4 )
          {
            if (std::find(cls_names.begin(), cls_names.end(), dt->d_name) != cls_names.end())
            {
              std::cout<<"find pretrain vector class = "<< dt->d_name<<std::endl;
              load_pre_text_embeddings(p.assign(path).append("/").append(dt->d_name), known_embeddings, cls_names);
            }
          } else
          {
            files.push_back(p.assign(path).append("/").append(dt->d_name));
          }
        }
      }

      closedir(d);
      if(!files.empty())
      {
        known_embeddings.push_back(load_pre_text_file(files));
      }
    }
    return true;
  }

  float cosine_similarity(const std::vector<float>& emb1,
                  const std::vector<float>& emb2) {
    if (emb1.size() != emb2.size())
    {
      std::cout<<"SIZE INCONSISTENT with two vectors, one is "<<emb1.size()<<", the other is "<< emb2.size()<<std::endl;
      return 0;
    }

    float dot = std::inner_product(emb1.begin(), emb1.end(), emb2.begin(), 0.0);
    float emb1_sum = std::inner_product(emb1.begin(), emb1.end(), emb1.begin(), 0.0);
    float emb2_sum = std::inner_product(emb2.begin(), emb2.end(), emb2.begin(), 0.0);
    dot /= std::max(std::sqrt(emb1_sum) * std::sqrt(emb2_sum),
                    std::numeric_limits<float>::epsilon());
    return dot;
  }

  std::vector<float> softmax(std::vector<float> &conf, float temperature)
  {
    for (auto &value : conf)
    {
      value *= temperature;
    }

    std::vector<float>::iterator biggest = std::max_element(conf.begin(),conf.end());
    float max_value = *biggest;
    
    float total = 0;
    for (auto x : conf)
    {
      total += exp(x-max_value);
    }

    std::vector<float> result;
    for (auto x : conf)
    {
      result.push_back(exp(x-max_value)/total);
    }
    return result;
  }

} // namespace optical_clip_detector

